package com.aikg.kgrag_java.service.impl;

import com.aikg.kgrag_java.pojo.entity.ConversationInfo;
import com.aikg.kgrag_java.pojo.entity.Prompts;
import com.aikg.kgrag_java.pojo.enums.RoleEnum;
import com.aikg.kgrag_java.pojo.request.LLMRequestBody;
import com.aikg.kgrag_java.service.ILLMService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @ClassName: QuestionRecommendPipeline
 * @Description: 问题推荐流程服务类
 * @Author: 李宗雨
 * @Date: 2025/1/9 16:59
 * @Version: 1.0
 **/
@Service
public class QuestionRecommendPipeline {

    @Autowired
    ILLMService llmService;

    /**
     * 问题推荐，按照上下文推荐k个问题
     * @param conversationInfo 问答对信息
     * @param k k个问题
     * @return k个问题
     */
    public List<String> questionRecommend(ConversationInfo conversationInfo, int k){
        if(conversationInfo.getChunks().isEmpty()){
            // 上下文为空则直接返回空
            return new ArrayList<>();
        }
        // 拼接判断问题和答案出处的prompt，并拼接问答对话数组
        LLMRequestBody llmRequestBody = new LLMRequestBody();
        llmRequestBody.addMessages(RoleEnum.USER, Prompts.getQuestionRecommendPrompt1(conversationInfo));

        // 正则表达式判断模型返回的结果中是否含有段落编号
        Pattern pattern = Pattern.compile("\\d+");
        Matcher matcher = pattern.matcher(llmService.chat(llmRequestBody));
        String numStr = "";
        if (matcher.find()) {
            numStr = matcher.group();
        }
        String prompt = "";
        if (numStr.length() == 1){
            // 得到的数字是个位数
            int num = Integer.parseInt(numStr);
            if (1 <= num && num <= conversationInfo.getChunks().size()){
                // 得到的数字在合理范围内
                prompt = Prompts.getQuestionRecommendPrompt2(conversationInfo.getChunks().get(num - 1).getPage_content(), k);
            }
        }
        // 若未得到正确的段落号则prompt拼接会不成功，为空则直接返回
        if(prompt.isEmpty()){
            return new ArrayList<>();
        }

        // 拼接问题推荐的prompt，并拼接问答对话数组
        llmRequestBody = new LLMRequestBody();
        llmRequestBody.addMessages(RoleEnum.USER,
                "上下文信息：包括实战和模拟环境演习在内的电子战训练；电子战装备的维护工作以保证装备的整洁好用；与参谋的其它成员共同参与军事决策过程。实践军事决策程序培养团队合作并建立电子战指挥官发出或是参谋长反馈的期望。包括通信情报和电子战资源和能力的协同演练；计划、发布和报道电子战资源的动向；在机动过程中，协调清晰的路线和护送要求以避免风险和延误。4-3. 电子战准备过程中，包括：更新对电子战运行的评估使得与情报运行保持一致；向情报助理参谋长和信号助理参谋长申请对联合限制频率列表和信号操作条例的更改和特殊条款的批复；全风险评估和相应的规避策略；领导赛博空间电磁行动工作团队；组织军事演习和人员配置包括：配置电子攻击请求模式和评估电子战行动的效果；开发电子战战场和机载控制授权程序；整合信息收集人员（情报助理参谋长和信号助理参谋长）；协调外部维修和重置电子战装备的支撑程序；启动电子防护程序用来对抗电磁干扰和敌军的干扰行动；开发标准操作流程；建立报告程序；实行战前核实和电子战装备的检查。避免电磁频谱的冲突4-4. 要想避免电磁频谱使用冲突，需要了解和任务需求。电子战指挥官要考虑装备的作用距离、定位精度和对友好频率以及受限频谱依赖程度，并给出建议信号操作条例和联合限制频率列表。信号操作条例包括呼叫信号、呼叫语音、频率分配、符号并分配给友军。关于信号操作条例和联合限制频率列表详细信息请参考ATP 6-02.70。\n" +
                        "\n" +
                        "根据提供的上下文信息和没有先验知识的原则，仅基于以下查询生成问题。\n" +
                        "你是一名教师/教授。你的任务是围绕上下文信息设置3个主题相关问题。这些问题的答案必须在上下文信息中找得到。仅输出三对问题和答案，其他多余信息不要输出。"
        );
        llmRequestBody.addMessages(RoleEnum.ASSISTANT,
                "问题:在电子战训练中，哪些方面是电子战指挥官需要关注的？\n" +
                        "答:电子战指挥官需要关注包括实战和模拟环境演习在内的电子战训练，电子战装备的维护工作以保证装备的整洁好用，以及与参谋的其他成员共同参与军事决策过程。\n" +
                        "\n" +
                        "问题:在电子战准备过程中，哪些工作是电子战指挥官需要完成的？\n" +
                        "答:电子战指挥官需要在准备过程中完成更新对电子战运行的评估使得与情报运行保持一致，向情报助理参谋长和信号助理参谋长申请对联合限制频率列表和信号操作条例的更改和特殊条款的批复，进行全风险评估和相应的规避策略，领导赛博空间电磁行动工作团队，组织军事演习和人员配置等。\n" +
                        "\n" +
                        "问题:如何避免电磁频谱使用冲突？\n" +
                        "答:为了避免电磁频谱使用冲突，电子战指挥官需要了解任务需求，考虑装备的作用距离、定位精度和对友好频率以及受限频谱依赖程度，并给出建议信号操作条例和联合限制频率列表。信号操作条例包括呼叫信号、呼叫语音、频率分配、符号并分配给友军。"
        );
        llmRequestBody.addMessages(RoleEnum.USER, prompt);

        // 模型问答返回问题推荐列表
        String questionsStr = llmService.chat(llmRequestBody);

        // 对模型的返回做处理得到k个问题
        String[] questions = questionsStr.split("\\n+");
        List<String> result = new ArrayList<>();
        for (int i = 0; i< questions.length;i+=2){
            int index = -1;
            if (questions[i].indexOf(":") == 2){
                // 找到了
                index = questions[i].indexOf(":");
            }
            else if (questions[i].indexOf("：") == 2){
                // 找到了
                index = questions[i].indexOf("：");
            }
            else{
                continue;
            }
            result.add(questions[i].substring(index + 1));
        }
        return result;
    }
}
